//
// Created by 46508 on 2024/5/26.
//

#include <VkDevice.hpp>
#include <iostream>
#include <sstream>
#include <vector>
#include <vulkan/vulkan.hpp>
#include <functional>


using EnumerateQueueFamilyIdxChecker = const std::function<bool(const vk::QueueFamilyProperties &, u32 queueFamilyIdx)>;


/// Helper function to convert memory property flags to a string
static std::string MemoryTypeToString(const vk::MemoryPropertyFlags flags) {
    std::stringstream result;
    if (vk::MemoryPropertyFlagBits::eDeviceLocal & flags) {
        result << "|本地设备";
    }
    if (vk::MemoryPropertyFlagBits::eHostVisible & flags) {
        result << "|主机可见";
    }
    if (vk::MemoryPropertyFlagBits::eHostCoherent & flags) {
        result << "|主机一致";
    }
    if (vk::MemoryPropertyFlagBits::eHostCached & flags) {
        result << "|主机缓存";
    }
    if (vk::MemoryPropertyFlagBits::eLazilyAllocated & flags) {
        result << "|惰性分配";
    }
    result << "|";
    return result.str();
}

/// 打印设备信息
static void PrintDeviceInfo(const vk::PhysicalDevice &device) {
    const auto properties = device.getProperties();
    const auto features = device.getFeatures();
    const auto memoryProperties = device.getMemoryProperties();
    std::cout << "设备ID: " << properties.deviceID << std::endl;
    std::cout << "设备类型: ";
    std::cout << vk::to_string(properties.deviceType) << std::endl;
    std::cout << "设备名称: " << properties.deviceName << std::endl;
    for (uint32_t i = 0; i < memoryProperties.memoryTypeCount; i++) {
        const auto memoryType = memoryProperties.memoryTypes[i];
        const auto index = memoryType.heapIndex;
        const auto memoryHeap = memoryProperties.memoryHeaps[index];
        const auto memoryTypeStr = MemoryTypeToString(memoryType.propertyFlags);
        std::cout << "显存类型: " << memoryTypeStr << " 大小:" << memoryHeap.size / (1024 * 1024) << "MB" << std::endl;
    }
    std::cout << "设备版本：" << properties.driverVersion << std::endl;
    std::cout << "几何着色器: " << (features.geometryShader ? "支持" : "不支持") << std::endl;
}

/// 获取设备评分
static u32 GetPhysicalDeviceScoring(const vk::PhysicalDevice &device) {
    uint32_t score = 0;
    const auto properties = device.getProperties();
    const auto features = device.getFeatures();
    const auto memoryProperties = device.getMemoryProperties();
    // 根据设备类型评分
    switch (properties.deviceType) {
        case vk::PhysicalDeviceType::eDiscreteGpu:
            score += 2048;
            break;
        case vk::PhysicalDeviceType::eIntegratedGpu:
            score += 1024;
            break;
        case vk::PhysicalDeviceType::eVirtualGpu:
            score += 512;
            break;
        case vk::PhysicalDeviceType::eCpu:
            score += 128;
            break;
        default:
            break;
    }
    // 根据最大图像维度评分
    score += properties.limits.maxImageDimension2D / 1024;
    // 根据内存大小评分
    uint64_t memorySize = 0;
    for (uint32_t i = 0; i < memoryProperties.memoryHeapCount; i++) {
        if (memoryProperties.memoryHeaps[i].flags & vk::MemoryHeapFlagBits::eDeviceLocal) {
            memorySize += memoryProperties.memoryHeaps[i].size;
        }
    }
    // 每64MB得1分
    score += static_cast<uint32_t>(memorySize / (1024 * 1024)) / 32;
    // 根据驱动程序版本评分
    score += properties.driverVersion / 100000;
    // 根据其他特性评分
    if (features.geometryShader) score += 128;
    return score;
}

/// 枚举队列族索引
static void EnumerateQueueFamilyIdx(const vk::PhysicalDevice &physicalDevice,
                                    EnumerateQueueFamilyIdxChecker &checker) {
    const auto properties = physicalDevice.getQueueFamilyProperties();
    for (int i = 0; i < properties.size(); ++i)
        if (!checker(properties[i], i)) return;
}

static auto CreateQueueCreateInfo(const u32 queueFamilyIdx,
                                  const f32 queuePriorities = 1,
                                  const u32 queueCount = 1) -> vk::DeviceQueueCreateInfo {
    auto createInfo = vk::DeviceQueueCreateInfo{};
    createInfo.setQueueFamilyIndex(queueFamilyIdx)
            .setQueueCount(queueCount)
            .setPQueuePriorities(&queuePriorities);
    return createInfo;
}


namespace NaturalEngine {
    VkDevice::VkDevice(const IDeviceExtensionsProvider &extensionsProvider,
                       const IDeviceSurfaceProvider &surfaceProvider
    ) {
        vk::ApplicationInfo appInfo{};
        appInfo.setApiVersion(VK_API_VERSION_1_3)
                .setPApplicationName(ENGINE_NAME)
                .setPEngineName(APP_NAME);
        vector<cstr> extensions{};
        extensionsProvider.GetExtensions(extensions);
        CreateInstance(appInfo, extensions);
        GetPhysicalDevice();
        surfaceProvider.CreateSurface(_instance, _surface);
        QueryQueueFamilyIdx();
        CreateLogicDevice();
        QueryQueue();
    }

    VkDevice::~VkDevice() {
        if (_logicDevice) _logicDevice.destroy();
        if (_surface) _instance.destroySurfaceKHR(_surface);
        if (_instance)_instance.destroy();
    }

    void VkDevice::CreateSharedSwapchainKHR(vk::SwapchainKHR &swapchain, vk::SwapchainCreateInfoKHR &createInfo) const {
        const auto &graphicsQueueFamilyIdx = _graphicsQueue.FamilyIdx;
        const auto &presentFamilyQueueIdx = _graphicsQueue.FamilyIdx;
        vector<u32> familyQueueIdxs{graphicsQueueFamilyIdx.value()};
        auto sharingMode = vk::SharingMode::eExclusive;
        if (graphicsQueueFamilyIdx.value() != presentFamilyQueueIdx.value()) {
            familyQueueIdxs.push_back(presentFamilyQueueIdx.value());
            sharingMode = vk::SharingMode::eConcurrent;
        }
        createInfo.setQueueFamilyIndices(familyQueueIdxs)
                .setSurface(_surface)
                .setImageSharingMode(sharingMode);
        swapchain = _logicDevice.createSharedSwapchainKHR(createInfo);
    }

    VkSwapChainInfo VkDevice::QuerySwapChainInfo(const i32 width, const i32 height, const u32 imageCount) const {
        VkSwapChainInfo info;
        for (const auto &surfaceFormat: _physicalDevice.getSurfaceFormatsKHR()) {
            if (vk::Format::eR8G8B8A8Srgb == surfaceFormat.format && vk::ColorSpaceKHR::eSrgbNonlinear == surfaceFormat.
                colorSpace) {
                info.Format = surfaceFormat;
            }
        }
        const auto capabilities = _physicalDevice.getSurfaceCapabilitiesKHR(_surface);

        const auto minImageExtent = capabilities.minImageExtent;
        const auto maxImageExtent = capabilities.maxImageExtent;
        info.ImageCount = std::clamp<u32>(imageCount, capabilities.minImageCount, capabilities.maxImageCount);
        info.ImageExtent.width = std::clamp<u32>(width, minImageExtent.width, maxImageExtent.width);
        info.ImageExtent.height = std::clamp<u32>(height, minImageExtent.height, minImageExtent.height);
        info.Transform = capabilities.currentTransform;

        const auto presents = _physicalDevice.getSurfacePresentModesKHR(_surface);
        info.Present = vk::PresentModeKHR::eFifo;
        for (const auto &present: presents) {
            if (vk::PresentModeKHR::eMailbox == present) {
                info.Present = present;
                break;
            }
        }
        return info;
    }

    void VkDevice::CreateInstance(const vk::ApplicationInfo &appInfo, const vector<cstr> &extensions) {
        // layer 层 嵌入到vk中用于调试
        // extension 扩展 包含功能
        vk::InstanceCreateInfo createInfo{};
        std::vector<const char *> layers = {};
        for (auto &layer: vk::enumerateInstanceLayerProperties()) {
#ifdef DEBUG
            if (std::strcmp("VK_LAYER_LUNARG_monitor", layer.layerName) == -1)
                layers.push_back("VK_LAYER_LUNARG_monitor");
            if (std::strcmp("VK_LAYER_KHRONOS_validation", layer.layerName) == -1)
                layers.push_back("VK_LAYER_KHRONOS_validation");
#endif
        }
        createInfo.setPApplicationInfo(&appInfo)
                .setPEnabledLayerNames(layers)
                .setPEnabledExtensionNames(extensions);
        _instance = createInstance(createInfo);
    }

    void VkDevice::GetPhysicalDevice() {
        auto devices = _instance.enumeratePhysicalDevices();
        vk::PhysicalDevice bestDevice;
        u32 highestScore = 0;
        for (const auto &device: devices) {
            PrintDeviceInfo(device);
            if (const u32 score = GetPhysicalDeviceScoring(device); score > highestScore) {
                highestScore = score;
                bestDevice = device;
            }
        }
        if (!bestDevice) {
            throw std::runtime_error("Failed to find a suitable GPU");
        }
        _physicalDevice = bestDevice;
    }

    void VkDevice::QueryQueueFamilyIdx() {
        std::optional<u32> graphicsQueueFamilyIdx{};
        std::optional<u32> presentQueueFamilyIdx{};
        EnumerateQueueFamilyIdx(_physicalDevice, [&](const vk::QueueFamilyProperties &property, u32 queueFamilyIdx) {
            if (property.queueFlags & vk::QueueFlagBits::eGraphics)
                graphicsQueueFamilyIdx = queueFamilyIdx;
            if (_physicalDevice.getSurfaceSupportKHR(queueFamilyIdx, _surface))
                presentQueueFamilyIdx = queueFamilyIdx;
            return presentQueueFamilyIdx.has_value() && graphicsQueueFamilyIdx.has_value();
        });
        _graphicsQueue.FamilyIdx = graphicsQueueFamilyIdx;
        _presentQueue.FamilyIdx = presentQueueFamilyIdx;
    }

    void VkDevice::CreateLogicDevice() {
        const auto &graphicsQueueFamilyIdx = _graphicsQueue.FamilyIdx.value();
        const auto &presentQueueFamilyIdx = _presentQueue.FamilyIdx.value();
        vector<vk::DeviceQueueCreateInfo> queueCreateInfos{};
        if (graphicsQueueFamilyIdx != presentQueueFamilyIdx) {
            queueCreateInfos.push_back(CreateQueueCreateInfo(graphicsQueueFamilyIdx));
            queueCreateInfos.push_back(CreateQueueCreateInfo(presentQueueFamilyIdx));
        } else {
            queueCreateInfos.push_back(CreateQueueCreateInfo(graphicsQueueFamilyIdx));
        }
        vk::DeviceCreateInfo deviceCreateInfo{};
        deviceCreateInfo.setQueueCreateInfos(queueCreateInfos);
        deviceCreateInfo.setQueueCreateInfoCount(queueCreateInfos.size());
        _logicDevice = _physicalDevice.createDevice(deviceCreateInfo);
    }

    void VkDevice::QueryQueue() {
        const auto graphicsQueueFamilyIdx = _graphicsQueue.FamilyIdx.value();
        const auto presentQueueFamilyIdx = _presentQueue.FamilyIdx.value();
        _graphicsQueue.Queue = _logicDevice.getQueue(graphicsQueueFamilyIdx, 0);
        _graphicsQueue.Queue = _logicDevice.getQueue(presentQueueFamilyIdx, 0);
    }
}
